import functools
import gzip
import struct
from io import BytesIO
from typing import List, Type, Dict, Optional
from emi.io import write_float_compress, read_float_compress, write_int_compress, write_value, \
    write_bytes_list, read_bytes_list, read_int_compress, write_dict
from emi.bar.storage.model import BaseDBModel
from emi.bar.db import BarDB, DayDB
from emi.barx.data import BarX
from emi.core.DBSource import DBSource
from emi.util import  TimeUtils
from emi.bar.data import Interval
import peewee
from datetime import datetime
import numpy as np

def to_datetime(yy_mm_dd:int,the_day_end = False)->datetime:
    year = int(yy_mm_dd / 10000 + 0.01)
    month = int((yy_mm_dd / 100 + 0.01)) % 100
    day = yy_mm_dd % 100
    if the_day_end:
        return datetime(year=year, month=month, day=day, hour=23, minute=59, second=59)
    else:
        return datetime(year=year, month=month, day=day, hour=0, minute=0, second=0)


def _make_detail_bytes(base: BarX, details: List[BarX]):
    output = BytesIO()
    bars = details
    hh_mm_ss_list = np.zeros(len(bars), int)
    volume_list = np.zeros(len(bars), float)
    open_list = np.zeros(len(bars), float)
    high_list = np.zeros(len(bars), float)
    low_list = np.zeros(len(bars), float)
    close_list = np.zeros(len(bars), float)
    amount_list = np.zeros(len(bars), float)
    extra_list = []
    item_data_list = []

    y_m_d, __h_m_s = TimeUtils.split_datetime(base.datetime)
    pre_h_m_s = __h_m_s
    for i in range(0, len(bars)):
        bar = bars[i]
        the_y_m_d, h_m_s = TimeUtils.split_datetime(bar.datetime)
        ##保证批量存储的数据是同一天和同一维度
        if not the_y_m_d == y_m_d:
            raise RuntimeError("pack_bars_as_model():bar的datetime不同一天")
        if pre_h_m_s > h_m_s:
            raise RuntimeError("pack_bars_as_model():存储的bars要升序排序")
        hh_mm_ss_list[i] = np.int16(h_m_s / 100)
        volume_list[i] = int(bar.volume)
        open_list[i] = bar.open
        close_list[i] = bar.close
        high_list[i] = bar.high
        low_list[i] = bar.low
        amount_list[i] = bar.amount
        pre_h_m_s = h_m_s
        _extra = bar.extra
        if _extra:
            extra_output = BytesIO()
            write_value(extra_output, _extra)
            extra_list.append(extra_output.getvalue())
        else:
            extra_list.append(None)

        _item_data = bar._item_map()
        if _item_data:
            item_output = BytesIO()
            write_value(item_output, _item_data)
            item_data_list.append(item_output.getvalue())
        else:
            item_data_list.append(None)

    write_int_compress(output, hh_mm_ss_list)
    write_int_compress(output, volume_list)
    write_float_compress(output, open_list, presicion=3)
    write_float_compress(output, close_list, presicion=3)
    write_float_compress(output, high_list, presicion=3)
    write_float_compress(output, low_list, presicion=3)
    write_float_compress(output, amount_list, presicion=3)
    write_bytes_list(output,extra_list)
    write_bytes_list(output,item_data_list)

    return gzip.compress(output.getvalue())


def _load_details(core_data, base: BarX)->List[BarX]:
    core_data = gzip.decompress(core_data)
    hh_mm_ss_list, offset = read_int_compress(core_data, 0)
    volume_list, offset = read_int_compress(core_data, offset)
    open_list, offset = read_float_compress(core_data, offset)
    close_list, offset = read_float_compress(core_data, offset)
    high_list, offset = read_float_compress(core_data, offset)
    low_list, offset = read_float_compress(core_data, offset)
    amount_list, offset = read_float_compress(core_data, offset)
    extra_list,offset = read_bytes_list(core_data,offset)
    item_data_list,offset = read_bytes_list(core_data,offset)

    bars = []
    today = base.datetime
    pre_close = base.pre_close
    price_adjust = base.price_adjust
    for i in range(0, len(hh_mm_ss_list)):
        hh_mm_ss = hh_mm_ss_list[i] * 100
        hour = int(hh_mm_ss / 10000 + 0.01)
        minute = int((hh_mm_ss / 100) % 100 + 0.01)
        second = int(hh_mm_ss % 100 + 0.01)
        bar = BarX(interval=base.interval, symbol=base.symbol,
                      datetime=today.replace(hour=hour, minute=minute, second=second))
        bar.open = open_list[i]
        bar.close = close_list[i]
        bar.high = high_list[i]
        bar.low = low_list[i]
        bar.volume = volume_list[i]
        bar.amount = amount_list[i]
        bar.pre_close = pre_close
        bar.price_adjust = price_adjust
        bar._extra_bytes = extra_list[i]
        bar._item_bytes = item_data_list[i]
        bars.append(bar)
    return bars


class BarXModel(BaseDBModel):
    """
    dimen_value:str = None  ## 维度值


    open: float = 0  ## 复权后的价格
    high: float = 0  ## 复权后的价格
    low: float = 0    ## 复权后的价格
    close: float = 0   ## 复权后的价格
    pre_close:float = 0 ## 上个交易日的收盘价。 一般情况下是没有的，需要执行BarStorage.update_pre_close()才有
    volume: float = 0
    amount:float = 0  ## 成交金额

    _extra_bytes:bytes = None
    _item_bytes:bytes = None
    """

    symbol: str = peewee.CharField(index=True,max_length=64)
    year_month_day:int = peewee.SmallIntegerField(index=True)
    interval: int = peewee.SmallIntegerField(index=True)
    dimen_value:str = peewee.CharField(null=True,max_length=512)

    base_bytes:bytes = peewee.BlobField(null=True)  ## open,high,low,close,pre_close,volume,amount
    extra_bytes: bytes = peewee.BlobField(null=True)
    item_bytes: bytes = peewee.BlobField(null=True)
    detail_bytes:bytes = peewee.BlobField(null=True)

    def to_dict(self):
        return self.__data__

    @classmethod
    def get_field(cls, detail: bool):
        db_model = cls
        if detail:
            return (
                db_model.symbol, db_model.year_month_day, db_model.interval, db_model.dimen_value,
                db_model.base_bytes, db_model.extra_bytes, db_model.item_bytes,
                db_model.detail_bytes)
        else:
            return (
                db_model.symbol, db_model.year_month_day, db_model.interval, db_model.dimen_value,
                db_model.base_bytes, db_model.extra_bytes, db_model.item_bytes,
            )

    @staticmethod
    def from_data(bar: BarX):
        db_bar = BarXModel()
        db_bar.symbol = bar.symbol
        db_bar.interval = bar.interval.value
        y_m_d, h_m_s = TimeUtils.split_datetime(bar.datetime)
        db_bar.year_month_day = y_m_d
        db_bar.dimen_value = bar.dimen_value

        output = BytesIO()
        ##压缩存储,保留3个精度。
        write_float_compress(output, [bar.open, bar.high, bar.low, bar.close, bar.pre_close], presicion=3)
        output.write(struct.pack("=qqd", int(bar.volume), int(bar.amount * 1000),bar.price_adjust))
        db_bar.base_bytes = output.getvalue()

        _extra = bar.extra
        if _extra:
            output = BytesIO()
            write_dict(output, _extra)
            db_bar.extra_bytes = output.getvalue()

        _item_data = bar._item_map()
        if _item_data:
            output = BytesIO()
            write_dict(output, _item_data)
            db_bar.item_bytes = output.getvalue()

        details =  bar.details
        if details:
            db_bar.detail_bytes = _make_detail_bytes(bar,details)

        return db_bar


    def to_bar(self)->BarX:
        year = int(self.year_month_day / 10000 + 0.01)
        month = int((self.year_month_day / 100 + 0.01)) % 100
        day = self.year_month_day % 100
        time = datetime(year=year, month=month, day=day, hour=15, minute=0, second=0)
        bar = BarX(interval = Interval(self.interval),
                   datetime = time,
                   symbol = self.symbol,
                   dimen_value=self.dimen_value
                   )
        base_bytes = self.base_bytes
        (bar.open, bar.high, bar.low, bar.close, bar.pre_close), offset = read_float_compress(base_bytes, 0)
        bar.volume, _amount,bar.price_adjust = struct.unpack_from("=qqd", base_bytes, offset)
        bar.amount = _amount / 1000

        if self.extra_bytes:
            bar._extra_bytes = self.extra_bytes

        if self.item_bytes:
            bar._item_bytes = self.item_bytes

        if self.detail_bytes:
            bar._lazy_load_fun = functools.partial(_load_details,self.detail_bytes,bar)

        return bar

class BarXStorage(BarDB):

    def __init__(self, source: DBSource, table_name=None):
        indexes_desc = (
            (("symbol", "year_month_day", "interval"), True),
        )
        if table_name is None:
            table_name = 'barx_db_default'
        self.source = source.init_peewee(table_name,BarXModel,indexes=indexes_desc,ignore_inited=True)

    def get_bars(self, symbol: str, interval: Interval = None, start: datetime = None,
                 end: datetime = None,with_detail = True) -> List[BarX]:
        if start is None:
            start = datetime(year=1970, month=1, day=1, hour=0, minute=0, second=0)
        if end is None:
            end = datetime.now()
        if interval is None:
            interval = self.get_default_interval()

        start_y_m_d, h_m_s = TimeUtils.split_datetime(start)
        end_y_m_d, h_m_s = TimeUtils.split_datetime(end)
        interval = interval.value
        with self.source.use_peewee() as db_model:
            fields = db_model.get_field(with_detail)
            if symbol is None:
                s = (
                    db_model.select(*fields)
                        .where(
                        (db_model.interval == interval)
                        & (db_model.year_month_day >= start_y_m_d)
                        & (db_model.year_month_day <= end_y_m_d)
                    ).order_by(db_model.year_month_day)
                )
                data = [db_bar.to_bar() for db_bar in s]
                return data
            s = (
                db_model.select(*fields)
                    .where(
                    (db_model.symbol == symbol)
                    & (db_model.interval == interval)
                    & (db_model.year_month_day >= start_y_m_d)
                    & (db_model.year_month_day <= end_y_m_d)
                )
                    .order_by(db_model.year_month_day)
            )
            data = [db_bar.to_bar() for db_bar in s]
            return data

    def get_bars_recently(self, symbol: str, time: datetime, count: int, interval: Interval = None, is_before=True,with_detail = True) -> List[BarX]:
        _y_m_d, h_m_s = TimeUtils.split_datetime(time)
        if interval is None:
            interval = self.get_default_interval()
        interval = interval.value

        with self.source.use_peewee() as db_model:
            fields = db_model.get_field(with_detail)
            if is_before:
                s = (
                    db_model.select(*fields)
                        .where(
                        (db_model.symbol == symbol)
                        & (db_model.interval == interval)
                        & (db_model.year_month_day <= _y_m_d)
                    )
                        .order_by(db_model.year_month_day.desc())
                        .limit(count)
                )
                data = [db_bar.to_bar() for db_bar in s]
                return list(reversed(data))
            else:
                s = (
                    db_model.select(*fields)
                        .where(
                        (db_model.symbol == symbol)
                        & (db_model.interval == interval)
                        & (db_model.year_month_day >= _y_m_d)
                    )
                        .order_by(db_model.year_month_day.asc())
                        .limit(count)
                )
                data = [db_bar.to_bar() for db_bar in s]
                return data

    def get_bars_at_day(self, time: datetime, interval: Interval = None,with_detail = True) -> Dict[str, BarX]:
        """
        获取某一天的所有数据:
        key : symbol。
        :interval 行情频率，为None时取默认值：get_default_interval()
        """
        _y_m_d, h_m_s = TimeUtils.split_datetime(time)
        if interval is None:
            interval = self.get_default_interval()
        interval = interval.value

        with self.source.use_peewee() as db_model:
            fields = db_model.get_field(with_detail)
            s = (
                db_model.select(*fields)
                    .where(
                    (db_model.interval == interval)
                    & (db_model.year_month_day == _y_m_d)
                )
            )
            dict_data = {}
            for db_bar in s:
                bar = db_bar.to_bar()
                dict_data[bar.symbol] = bar
            return dict_data

    def get_newest_bar(self, symbol: str = None, interval=Interval.DAILY,with_detail:bool = True) -> Optional[BarX]:
        with self.source.use_peewee() as db_model:
            fields = db_model.get_field(with_detail)
            interval = interval.value
            s = (
                db_model.select(*fields)
                    .where(
                    (db_model.interval == interval)
                    & (db_model.symbol == symbol)
                )
                    .order_by(db_model.year_month_day.desc())
                    .first()
            )
            if s:
                return s.to_bar()
            return None

    def get_newest_datetime(self, symbol: str = None, interval=Interval.DAILY) -> Optional[datetime]:
        with self.source.use_peewee() as db_model:
            if interval:
                interval = interval.value
                if symbol is None:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.interval == interval)
                        )
                            .order_by(db_model.year_month_day.desc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
                else:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.interval == interval)
                            & (db_model.symbol == symbol)
                        )
                            .order_by(db_model.year_month_day.desc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
            else:
                if symbol is None:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .order_by(db_model.year_month_day.desc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
                else:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.symbol == symbol)
                        )
                            .order_by(db_model.year_month_day.desc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
            return None


    def get_oldest_datetime(self, symbol: str = None, interval=Interval.DAILY) -> Optional[datetime]:
        with self.source.use_peewee() as db_model:
            if interval:
                interval = interval.value
                if symbol is None:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.interval == interval)
                        )
                            .order_by(db_model.year_month_day.asc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
                else:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.interval == interval)
                            & (db_model.symbol == symbol)
                        )
                            .order_by(db_model.year_month_day.asc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
            else:
                if symbol is None:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .order_by(db_model.year_month_day.asc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
                else:
                    s = (
                        db_model.select(db_model.year_month_day)
                            .where(
                            (db_model.symbol == symbol)
                        )
                            .order_by(db_model.year_month_day.asc())
                            .first()
                    )
                    if s:
                        return to_datetime(s.year_month_day, the_day_end=True)
            return None

    def get_day_list(self, symbol: str = None, interval: Interval = None, start: datetime = None,
                     end: datetime = None) -> List[datetime]:
        """
        param:
            symbol 允许为None
        """
        year_month_day_list = self._load_day_list_with_yy_mm_sss(symbol, interval, start, end)
        return [to_datetime(y_m_d, the_day_end=True) for y_m_d in year_month_day_list]

    def get_symbol_list(self) -> List[str]:
        with self.source.use_peewee() as db_model:
            s = (db_model.select(db_model.symbol).distinct())
            return [v.symbol for v in s]

    def _load_day_list_with_yy_mm_sss(self, symbol: str = None, interval: Interval = None, start: datetime = None,
                                      end: datetime = None) -> List:
        """"
        加载股票symbol的行情日
        """
        begin = 1
        finished = 99999999
        if not start is None:
            _y_m_d, _h_m_s = TimeUtils.split_datetime(start)
            begin = _y_m_d
        if not end is None:
            _y_m_d, _h_m_s = TimeUtils.split_datetime(end)
            finished = _y_m_d

        with self.source.use_peewee() as db_model:
            if interval is None:
                if symbol is None:
                    s = (db_model.select(db_model.year_month_day).where(
                        ##(db_model.interval == interval)
                        (db_model.year_month_day >= begin)
                        & (db_model.year_month_day <= finished)
                    ).order_by(db_model.year_month_day.asc()).distinct())
                else:
                    s = (
                        db_model.select(db_model.year_month_day).where(
                            # (db_model.interval == interval)
                            (db_model.symbol == symbol)
                            & (db_model.year_month_day > begin)
                            & (db_model.year_month_day < finished)
                        ).order_by(db_model.year_month_day.asc()).distinct()
                    )
            else:
                interval = interval.value
                if symbol is None:
                    s = (db_model.select(db_model.year_month_day).where(
                        (db_model.year_month_day >= begin)
                        & (db_model.interval == interval)
                        & (db_model.year_month_day <= finished)
                    ).order_by(db_model.year_month_day.asc()).distinct())
                else:
                    s = (
                        db_model.select(db_model.year_month_day).where(
                            # (db_model.interval == interval)
                            (db_model.symbol == symbol)
                            & (db_model.interval == interval)
                            & (db_model.year_month_day > begin)
                            & (db_model.year_month_day < finished)
                        ).order_by(db_model.year_month_day.asc()).distinct()
                    )
            year_month_day_list = [db_bar.year_month_day for db_bar in s]
        return year_month_day_list

    def size(self, symbol: str = None) -> int:
        with self.source.use_peewee() as db_model:
            if symbol is None:
                return db_model.select().count()
            else:
                return db_model.select().where(db_model.symbol == symbol).count()

    def save_bars(self, bars: List[BarX]):
        """
        保存Bar数据。
        """
        ds = [BarXModel.from_data(i) for i in bars]
        with self.source.use_peewee() as db_model:
            db_model.save_all(ds)

    def clear_at_interval(self, interval: Interval):
        interval = interval.value
        with self.source.use_peewee() as db_model:
            db_model.delete().where(
                (db_model.interval == interval)
            ).execute()

    def clear_before(self, time: datetime):
        with self.source.use_peewee() as db_model:
            y_m_d, h_m_s = TimeUtils.split_datetime(time)
            db_model.delete().where(
                (db_model.year_month_day <= y_m_d)
            ).execute()

    def clear(self, symbol: str = None):
        with self.source.use_peewee() as db_model:
            if not symbol is None:
                db_model.delete().where(
                    (db_model.symbol == symbol)
                ).execute()
            else:
                ##情况表
                db_model.drop_table()
                db_model.create_table()


